package com.qzdatasoft.mcp.server.demo.prompt;

import com.qzdatasoft.mcp.server.common.AbstractPromptProvider;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.spec.McpSchema;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.function.BiFunction;

@Service
public class CalculatorPromptProvider extends AbstractPromptProvider {
    @Override
    protected McpSchema.Prompt getMeta() {
        return new McpSchema.Prompt(
                "calculator",
                "A simple calculator that can perform addition and subtraction.",
                List.of(
                        new McpSchema.PromptArgument("a", "first number", true),
                        new McpSchema.PromptArgument("b", "second number", true)
                )
        );
    }

    @Override
    protected BiFunction<McpSyncServerExchange, McpSchema.GetPromptRequest, McpSchema.GetPromptResult> handler() {
        return  (exchange, getPromptRequest) -> {
            String a = (String) getPromptRequest.arguments().get("a");
            String b = (String) getPromptRequest.arguments().get("b");
            var userMessage = new McpSchema.PromptMessage(
                    McpSchema.Role.USER,
                    new McpSchema.TextContent("The sum of " + a + " and " + b + " is " + (Integer.parseInt(a) + Integer.parseInt(b)) + ".")
            );
            return new McpSchema.GetPromptResult("A personalized greeting message", List.of(userMessage));
        };
    }


}
